Source code for hysop.operator.transpose

# Copyright (c) HySoP 2011-2024
#
# This file is part of HySoP software.
# See "https://particle_methods.gricad-pages.univ-grenoble-alpes.fr/hysop-doc/"
# for further info.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.


"""
@file transpose.py
Transposition related operators.
"""

from hysop.constants import Implementation, Backend, DirectionLabels
from hysop.tools.htypes import check_instance, to_tuple
from hysop.tools.decorators import debug
from hysop.tools.transposition_states import TranspositionState
from hysop.fields.continuous_field import Field, ScalarField

from hysop.topology.cartesian_descriptor import CartesianTopologyDescriptors
from hysop.core.graph.node_generator import ComputationalGraphNodeGenerator
from hysop.core.graph.computational_operator import ComputationalGraphNode

from hysop.testsenv import __HAS_OPENCL_BACKEND__


[docs] class TranspositionNotImplementedError(NotImplementedError): """ Error raised for unimplemented transposition operator configurations. """ pass
[docs] class Transpose(ComputationalGraphNodeGenerator): """ Operator generator for inplace and out of place field transposition and permutations in general. Available implementations are: *python (numpy based transpose, n-dimensional) *opencl (opencl codegen based transpose up to 16D) Implementations handle only one field at a time but this graph node generator generates an operator per supplied field, possibly on different implementation backends. See hysop.operator.base.transpose_operator.TransposeOperatorBase for operator backend implementation interface. """
[docs] @classmethod def implementations(cls): from hysop.backend.host.python.operator.transpose import PythonTranspose _implementations = { Implementation.PYTHON: PythonTranspose, } if __HAS_OPENCL_BACKEND__: from hysop.backend.device.opencl.operator.transpose import OpenClTranspose _implementations.update( { Implementation.OPENCL: OpenClTranspose, } ) return _implementations
[docs] @classmethod def default_implementation(cls): msg = "Transpose has no default implementation, " msg += "implementation should match the discrete field topology backend." raise RuntimeError(msg)
@debug def __new__( cls, fields, variables, axes, output_fields=None, implementation=None, name=None, base_kwds=None, **kwds, ): base_kwds = {} if (base_kwds is None) else {} return super().__new__( cls, name=name, candidate_input_tensors=None, candidate_output_tensors=None, **base_kwds, ) @debug def __init__( self, fields, variables, axes, output_fields=None, implementation=None, name=None, base_kwds=None, **kwds, ): """ Initialize a Transpose operator generator operating on CartesianTopology topologies. Parameters ---------- fields: Field, list or tuple of Fields Input continuous fields to be transposed, at least 2D. All fields should have the same dimension. output_fields: Field, list or tuple of Fields, optional Output continuous fields where the results are stored. Transposed shapes should match the input. By default output_fields are the same as input_fields resulting in inplace transpositions. Input and output are matched by order int list/tuple. variables: dict Dictionary of fields as keys and CartesianTopologyDescriptors as values. axes: tuple of ints, or array like of tuples, or dict of (tuple, TranspositionState). Permutation of axes in numpy notations (as a tuple of ints). Axe dim-1 is the contiguous axe, axe 0 has the greatest stride in memory. It can also be an array like of axes and the operator implementation will choose the most suitable axe permutation scheme in this list. If a dictionnary is provided, this gives the operator implementation to choose the most suitable axe permutation scheme with the knowledge of the target transposition state as well. All fields on the same backend will perform the same permutation. Operator chosen permutation can be retrieved by using generated operator 'axes' attribute. implementation: Implementation, optional, defaults to None target implementation, should be contained in available_implementations(). If implementation is set and topology does not match backend, RuntimeError will be raised on _generate. If None, implementation will be set according to topologies backend, different implementations may be choosen for different Fields if defined on different backends. name: string prefix for generated operator names base_kwds: dict, optional, defaults to None Base class keywords arguments. If None, an empty dict will be passed. kwds: keywords arguments that will be passed towards implementation transpose operator __init__. Notes ----- Out of place transpose will always be faster to process. The only exception to this rule may be 2D square matrices. Inplace transposition may request a temporary buffer because not all implementations may support inplace transposition. * Example for axes: - In 4D, axes=(0,1,2,3) means no permutations at all, this will raise ValueError. - In 2D, axes=(1,0) is a transposition. - In 3D, axes=(0,2,1) is a transposition in the XY-plane. * About dimensions: - No limit for the numpy CPU transposition implementation. - Max 16D for the OpenCL implementation. A Transpose operator implementation should support the TransposeOperatorBase interface (see hysop.operator.base.transpose_operator.TransposeOperatorBase). This ComputationalGraphNodeFrontend will generate a operator for each input and output ScalarField pair. All implementations should raise TranspositionNotImplementedError is the user supplied parameters leads to unimplemented or unsupported transposition features. """ input_fields = to_tuple(fields) assert None not in input_fields if output_fields is not None: output_fields = to_tuple(output_fields) output_fields = tuple( ofield if (ofield is not None) else ifield for (ifield, ofield) in zip(input_fields, output_fields) ) else: output_fields = tuple(ifield for ifield in input_fields) check_instance(input_fields, tuple, values=Field) check_instance(output_fields, tuple, values=Field, size=len(input_fields)) candidate_input_tensors = tuple(filter(lambda x: x.is_tensor, input_fields)) candidate_output_tensors = tuple(filter(lambda x: x.is_tensor, output_fields)) base_kwds = {} if (base_kwds is None) else {} super().__init__( name=name, candidate_input_tensors=candidate_input_tensors, candidate_output_tensors=candidate_output_tensors, **base_kwds, ) # expand tensors ifields, ofields = (), () for ifield, ofield in zip(input_fields, output_fields): msg = "Input and output field shape mismatch, got field {} of shape {} " msg += "and field {} of shape {}." if ifield.is_tensor ^ ofield.is_tensor: if ifield.is_tensor: msg = msg.format( ifield.short_description(), ifield.shape, ofield.short_description(), "(1,)", ) else: msg = msg.format( ifield.short_description(), "(1,)", ofield.short_description(), ofield.shape, ) raise RuntimeError(msg) if ifield.is_tensor and ofield.is_tensor and (ifield.shape != ofield.shape): msg = msg.format( ifield.short_description(), ifield.shape, ofield.short_description(), ofield.shape, ) raise RuntimeError(msg) ifields += ifield.fields ofields += ofield.fields input_fields = ifields output_fields = ofields check_instance(input_fields, tuple, values=ScalarField) check_instance(output_fields, tuple, values=ScalarField, size=len(input_fields)) check_instance(variables, dict, keys=Field, values=CartesianTopologyDescriptors) check_instance(base_kwds, dict, keys=str) check_instance(name, str, allow_none=True) fields = set(input_fields).union(output_fields) vfields = {f for tfield in variables.keys() for f in tfield.fields} if not fields: raise ValueError("fields are empty.") if not vfields: raise ValueError("variables are empty.") if not axes: raise ValueError("axes are empty.") if fields != vfields: if fields.difference(vfields): missing = tuple(field.name for field in fields - vfields) msg = "Missing fields in variables parameter: {}" msg = msg.format(", ".join(missing)) else: missing = tuple(field.name for field in vfields - fields) msg = "Too many fields present in variables keys: {}" msg = msg.format(", ".join(missing)) raise ValueError(msg) dim = input_fields[0].domain.dim if dim < 2: msg = "input_field dimension should be at least 2 but {} is a {}D field." msg = msg.format(input_fields[0].name, dim) raise ValueError(msg) for input_field, output_field in zip(input_fields, output_fields): in_topo_descriptor = ComputationalGraphNode.get_topo_descriptor( variables, input_field ) out_topo_descriptor = ComputationalGraphNode.get_topo_descriptor( variables, output_field ) if input_field.domain != output_field.domain: msg = "input_field {} and output_field {} do not share the same domain." msg.format(input_field.name, output_field.name) raise ValueError(msg) idim = input_field.domain.dim if idim != dim: msg = "input_field {} is of dimension {} and does not match first " msg += "input_field {} dimension {}." msg.format(input_field, idim, input_fields[0].name, dim) raise ValueError(msg) candidate_axes = self._check_axes(axes, dim) self.input_fields = input_fields self.output_fields = output_fields self.variables = variables self.candidate_axes = candidate_axes self.implementation = implementation self.kwds = kwds @staticmethod def _check_axes(axes, dim): msg = "Unkown type for axes parameter." if isinstance(axes, dict): candidate_axes = axes elif isinstance(axes, (list, tuple, set, frozenset)): if isinstance(axes[0], int): msg = "Axes should be ordered but got a set." assert not isinstance(axes, (set, frozenset)), msg candidate_axes = (tuple(axes),) elif isinstance(axes[0], (list, tuple)): candidate_axes = tuple(tuple(ax) for ax in axes) else: raise TypeError(msg) candidate_axes = dict(zip(candidate_axes, (None,) * len(candidate_axes))) else: raise TypeError(msg) del axes check_instance(candidate_axes, dict, keys=tuple) for axes, target_tstate in candidate_axes.items(): check_instance(axes, tuple, values=int) check_instance(target_tstate, TranspositionState[dim], allow_none=True) if len(axes) != dim: msg = "All axes should have the dimension of the transposed fields {} " msg += "but got axes={} which is of dimension {}." msg = msg.format(dim, axes, len(axes)) raise ValueError(msg) if set(axes) != set(range(dim)): msg = f"Axes should describe a valid pemutation, got axes={axes}." raise ValueError(msg) if axes == tuple(range(dim)): msg = "Given axes is identity (no permutation)." raise ValueError(msg) return candidate_axes def _get_op_and_check_implementation(self, src_topo, dst_topo): backend = getattr(src_topo, "backend", None) or getattr( dst_topo, "backend", None ) if backend is None: implementation = self.implementation if implementation is None: msg = ( "Source and destination topology descriptors do not expose backend " ) msg += ( "attribute and self.implementation was not set, cannot determine " ) msg += "the target transposition operator." raise ValueError(msg) if implementation not in self.implementations(): msg = f"Unknown transposition implementation {implementation}." raise ValueError(msg) op_cls = self.implementations()[implementation] else: def check_impl(impl, allowed): if impl not in allowed: msg = "Specified transposition implementation {} but this is not a valid " msg += "implementation for fields defined on backend of kind {}." msg = msg.format(impl, backend.kind) raise ValueError(msg) if backend.kind == Backend.OPENCL: implementation = self.implementation or Implementation.OPENCL check_impl(implementation, (Implementation.OPENCL,)) op_cls = self.implementations()[implementation] elif backend.kind == Backend.HOST: implementation = self.implementation or Implementation.PYTHON check_impl(implementation, (Implementation.PYTHON,)) op_cls = self.implementations()[implementation] else: assert implementation not in self.implementations() msg = f"Unknown transposition implementation {implementation}." raise ValueError(msg) from hysop.operator.base.transpose_operator import TransposeOperatorBase if not issubclass(op_cls, TransposeOperatorBase): msg = "Class {} does not inherit from the transpose operator interface " msg += "({}). This is an implementation error." msg = msg.format(op_cls, TransposeOperatorBase) raise TypeError(msg) if not issubclass(op_cls, ComputationalGraphNode): msg = ( "Class {} does not inherit from the computational graph node interface" ) msg += "({}). This is an implementation error." msg = msg.format(op_cls, ComputationalGraphNode) raise TypeError(msg) return op_cls @debug def _generate(self): nodes = [] for ifield, ofield in zip(self.input_fields, self.output_fields): src_topo = ComputationalGraphNode.get_topo_descriptor( self.variables, ifield ) dst_topo = ComputationalGraphNode.get_topo_descriptor( self.variables, ofield ) kwds = self.kwds.copy() TransposeOp = self._get_op_and_check_implementation(src_topo, dst_topo) axes = TransposeOp.get_preferred_axes( src_topo, dst_topo, self.candidate_axes ) variables = {ifield: src_topo} variables[ofield] = dst_topo # instantiate operator node = TransposeOp( input_field=ifield, output_field=ofield, variables=variables, axes=axes, **kwds, ) nodes.append(node) return nodes